from fastapi import APIRouter, Depends, HTTPException, UploadFile, File
from fastapi.responses import FileResponse
from sqlalchemy.orm import Session
from typing import List
from pydantic import BaseModel
from datetime import datetime
import os
import uuid
import aiofiles

from ..database import get_db
from ..models.user import User
from ..models.document import Document
from ..models.knowledge_base import KnowledgeBase
from ..auth.auth import get_current_active_user
from ..config import settings

router = APIRouter()

class DocumentResponse(BaseModel):
    id: int
    filename: str
    original_filename: str
    file_size: int
    file_type: str
    loader_type: str
    splitter_type: str
    chunk_count: int
    is_processed: bool
    is_vectorized: bool
    created_at: datetime
    knowledge_base_id: int
    
    class Config:
        from_attributes = True

def get_file_type(filename: str) -> str:
    ext = filename.lower().split('.')[-1]
    return ext

def is_allowed_file(filename: str) -> bool:
    return get_file_type(filename) in settings.ALLOWED_EXTENSIONS

@router.post("/upload/{kb_id}", response_model=DocumentResponse)
async def upload_document(
    kb_id: int,
    file: UploadFile = File(...),
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db)
):
    kb = db.query(KnowledgeBase).filter(
        KnowledgeBase.id == kb_id,
        KnowledgeBase.owner_id == current_user.id
    ).first()
    if not kb:
        raise HTTPException(status_code=404, detail="知识库不存在")
    
    if not is_allowed_file(file.filename):
        raise HTTPException(status_code=400, detail="不支持的文件类型")
    
    content = await file.read()
    if len(content) > settings.MAX_FILE_SIZE:
        raise HTTPException(status_code=400, detail="文件大小超出限制")
    
    file_ext = get_file_type(file.filename)
    unique_filename = f"{uuid.uuid4()}.{file_ext}"
    file_path = os.path.join(settings.UPLOAD_DIR, unique_filename)
    
    async with aiofiles.open(file_path, 'wb') as f:
        await f.write(content)
    
    document = Document(
        filename=unique_filename,
        original_filename=file.filename,
        file_path=file_path,
        file_size=len(content),
        file_type=file_ext,
        loader_type="auto",
        splitter_type="recursive",
        knowledge_base_id=kb_id
    )
    
    db.add(document)
    db.commit()
    db.refresh(document)
    
    return document

@router.get("/{kb_id}/documents", response_model=List[DocumentResponse])
async def list_documents(
    kb_id: int,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db)
):
    kb = db.query(KnowledgeBase).filter(
        KnowledgeBase.id == kb_id,
        KnowledgeBase.owner_id == current_user.id
    ).first()
    if not kb:
        raise HTTPException(status_code=404, detail="知识库不存在")
    
    documents = db.query(Document).filter(
        Document.knowledge_base_id == kb_id
    ).all()
    
    return documents

@router.delete("/{doc_id}")
async def delete_document(
    doc_id: int,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db)
):
    document = db.query(Document).join(KnowledgeBase).filter(
        Document.id == doc_id,
        KnowledgeBase.owner_id == current_user.id
    ).first()
    
    if not document:
        raise HTTPException(status_code=404, detail="文档不存在")
    
    if os.path.exists(document.file_path):
        os.remove(document.file_path)
    
    db.delete(document)
    db.commit()
    
    return {"message": "文档删除成功"} 